#coding=utf-8

# train based on yolov5su.pt with 200 Epochs

import cv2
import time 
from ultralytics import YOLO, solutions

#trt_model = YOLO("yolo11n.engine", task="detect")
trt_model = YOLO("CNR11n150E.engine", task="detect")

# 初始化目标裁剪器对象
# 参数说明：
# - show: 是否显示处理后的输出画面
# - model: 使用的目标检测模型路径
# - classes: 要裁剪的特定类别，这里[0, 1]通常对应于COCO数据集中的person和bicycle
# - conf: 可选，目标置信度阈值，默认为0.25
# - crop_dir: 可选，裁剪结果保存目录
cropper = solutions.ObjectCropper(
    # show=True,  # 显示输出结果
    model=trt_model,  # 用于目标裁剪的模型
    classes=[0, 1],  # 裁剪特定类别（person和bicycle）
    conf=0.7,  # 可调整目标的置信度阈值
    # crop_dir="cropped-detections",  # 设置裁剪检测结果的目录名称
)

# 从RTSP摄像头读取视频流
# RTSP URL格式: rtsp://username:password@ip_address:port/stream_path
# 请根据实际摄像头配置修改以下URL
rtsp_url = "rtsp://admin:a1234567@172.17.10.15:554/LiveMedia/ch1/Media1/trackID=1"
# 子码流
#rtsp_url = "rtsp://admin:a1234567@172.17.10.15:554/LiveMedia/ch1/Media2/trackID=1"
# rtsp_url = "gyl.mp4"
cap = cv2.VideoCapture(rtsp_url)
# 或增大缓冲区（增加延迟但减少丢包）
cap.set(cv2.CAP_PROP_BUFFERSIZE, 10)  # 较大缓冲区
# 确保RTSP流能够成功打开
assert cap.isOpened(), "Error connecting to RTSP stream"


last_time = time.time()
TARGET_INTERVAL = 0.1  # 100ms间隔

# 处理视频流
# 循环读取RTSP流的每一帧直到流关闭或发生错误
while cap.isOpened():
    current_time = time.time()
    if current_time - last_time >= TARGET_INTERVAL:
        # 读取一帧视频
        # success: 布尔值，表示是否成功读取帧
        # im0: 读取到的图像帧
        success, im0 = cap.read()

        # 如果无法读取帧（流结束或出错），则退出循环
        if not success:
            print("无法读取视频帧")
            #break
            cap.grab()

        else:
            # 将图像帧传递给裁剪器进行处理
            # 返回检测和裁剪的结果
            results = cropper(im0)

            # 可选：打印检测结果
            print(results)  # 访问输出结果

    time.sleep(0)

# 释放视频流资源
cap.release()
